#include <assert.h>
#include <string.h>
#include <math.h>
#include <random>

#include "mex.h"

#define VRBSLVL 0

#define INCLOGLKLHD(logpsbvi, tWZnt_d_ptr) \
{\
tsum = 0.0;\
cm1_dr_ptr = cm1_dr;\
cm2_dr_ptr = cm2_dr;\
half_invlambda_dr_ptr = half_invlambda_dr;\
for (int r = 0; r < Lm1; ++r) {\
	for (int d = 0; d < D; ++d) {\
		delta = cm1_dr_ptr[d]+cm2_dr_ptr[d]*((tWZnt_d_ptr)[d]+WZnt_d[d]);\
		tsum += delta*delta*half_invlambda_dr_ptr[d] + delta;\
	}\
	cm1_dr_ptr += D;\
	cm2_dr_ptr += D;\
	half_invlambda_dr_ptr += D;\
}\
(logpsbvi) -= tsum;\
}

int from_logps(double* const logps, const int N, const int algtype);

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
	/* Check inputs & outputs */
	if ((nrhs != 9 && nrhs != 10) || nlhs != 0)
		mexErrMsgTxt("Incorrect number of inputs/outputs!");
	for (int i = 0; i < nrhs; ++i)
		if (mxIsSparse(prhs[i]))
			mexErrMsgTxt("Inputs must be full matrices!");
	const bool isreal = nrhs==10;

	/* Input */
	double* const Zn_k = mxGetPr(prhs[0]);
	const double* const logmbyNmm_k = mxGetPr(prhs[1]);
	const double* const W_dk = mxGetPr(prhs[2]);
	const double* const cm1_dr = mxGetPr(prhs[3]);
	const double* const cm2_dr = mxGetPr(prhs[4]);
	const double* const half_invlambda_dr = mxGetPr(prhs[5]);
	double* const WZnt_d = mxGetPr(prhs[6]);
	int blocksz = (int)mxGetScalar(prhs[7]);
	const int algtype = (int)mxGetScalar(prhs[8]);
	const double* rvals = NULL;
	if (isreal) {
		rvals = mxGetPr(prhs[9]);
	}
	assert(algtype==1 || algtype==2);
	
	/* Compute Sizes */
	const int K = mxGetNumberOfElements(prhs[0]);
	const int D = mxGetNumberOfElements(prhs[6]);
	const int Lm1 = mxGetN(prhs[3]);
	int p2k = 1<<blocksz;
	
	const double *cm1_dr_ptr = cm1_dr, *cm2_dr_ptr = cm2_dr,
		*W_dk_ptr = W_dk, *tW_dk_ptr = W_dk,
		*half_invlambda_dr_ptr = half_invlambda_dr;
	double tsum, delta;

	if (! isreal) { // block sampler
		double* const logpsbv_p2k = (double*)mxMalloc(p2k*sizeof(double));
		double* const tWZnt_dp2k = (double*)mxMalloc(D*p2k*sizeof(double));

		double *tWZnt_dp2k_ptr = tWZnt_dp2k;
	
		if (VRBSLVL > 0) printf("[");

		int kend = 0;
		for (int k = 0; k < K; k = kend) {
			kend = k + blocksz;
			if (kend > K) {
				kend = K;
				blocksz = K - k;
				p2k = 1<<blocksz;
			}

			/* log-prior */
			logpsbv_p2k[0] = 0;
			for (int i = 2, j = kend-2; j >= k; i <<= 1, --j) {
				logpsbv_p2k[i] = logmbyNmm_k[j];
			}
			for (int pivot = 4, i = 6, n = 1; i < p2k-1;) {
				double pbase = logpsbv_p2k[pivot];
				for (int j = 0; j < n; ++j, i += 2) {
					logpsbv_p2k[i] = pbase + logpsbv_p2k[i-pivot];
				}
				i += 2;
				n = pivot - 1;
				pivot <<= 1;
			}
			const double logprior_end = logmbyNmm_k[kend-1];
			for (int i = 0; i < p2k; i += 2) {
				logpsbv_p2k[i+1] = logpsbv_p2k[i] + logprior_end;
			}
		
			/* log-likelihood */
			tW_dk_ptr = W_dk_ptr;
			for (int kk = k; kk < kend; ++kk) {
				for (int d = 0; d < D; ++d) {
					WZnt_d[d] -= (*tW_dk_ptr)*Zn_k[kk];
					tW_dk_ptr++;
				}
			}

			memset(tWZnt_dp2k, 0, D*p2k*sizeof(double));
			INCLOGLKLHD(logpsbv_p2k[0], tWZnt_dp2k);

			tW_dk_ptr = W_dk_ptr + D*(blocksz-2);
			tWZnt_dp2k_ptr = tWZnt_dp2k + (D<<1);
			for (int i = 2, j = kend-2; j >= k; i <<= 1, --j) {
				for (int d = 0; d < D; ++d) {
					tWZnt_dp2k_ptr[d] += tW_dk_ptr[d];
				}
				INCLOGLKLHD(logpsbv_p2k[i], tWZnt_dp2k_ptr);
			
				tWZnt_dp2k_ptr += D*i;
				tW_dk_ptr -= D;
			}

			tWZnt_dp2k_ptr = tWZnt_dp2k + D*6;
			for (int pivot = 4, i = 6, n = 1; i < p2k-1;) {
				const double *base_ptr = tWZnt_dp2k + D*pivot;
				const double *tptr = tWZnt_dp2k + D*(i-pivot);
				for (int j = 0; j < n; ++j, i += 2) {
					for (int d = 0; d < D; ++d) {
						tWZnt_dp2k_ptr[d] = base_ptr[d] + tptr[d];
					}
					INCLOGLKLHD(logpsbv_p2k[i], tWZnt_dp2k_ptr);

					tWZnt_dp2k_ptr += (D<<1);
					tptr += (D<<1);
				}
				i += 2;
				tWZnt_dp2k_ptr += (D<<1);
				n = pivot - 1;
				pivot <<= 1;
			}

			tW_dk_ptr = W_dk_ptr + D*(blocksz-1);
			tWZnt_dp2k_ptr = tWZnt_dp2k;
			for (int i = 0; i < p2k; i += 2) {
				for (int d = 0; d < D; ++d) {
					tWZnt_dp2k_ptr[D+d] = tWZnt_dp2k_ptr[d] + tW_dk_ptr[d];
				}
				tWZnt_dp2k_ptr += D;
				INCLOGLKLHD(logpsbv_p2k[i+1], tWZnt_dp2k_ptr);

				tWZnt_dp2k_ptr += D;
			}

			if (VRBSLVL > 1) {
				for (int i = 0; i < p2k; ++i) {
					printf("%.2f ", logpsbv_p2k[i]);
				}
			}

			int pid = from_logps(logpsbv_p2k, p2k, algtype);
			assert(pid >= 0);
			for (int kk = kend-1; kk >= k; --kk) {
				Zn_k[kk] = pid & 1;
				pid >>= 1;
			}
			
			for (int kk = k; kk < kend; ++kk) {
				for (int d = 0; d < D; ++d) {
					WZnt_d[d] += (*W_dk_ptr)*Zn_k[kk];
					W_dk_ptr++;
				}
			}

			if (VRBSLVL > 0) {
				for (int kk = k; kk < kend; ++kk) {
					printf("%d", (int)Zn_k[kk]);
				}
				printf(", ");
			}
		}
		if (VRBSLVL > 0) printf("\b\b]\n");

		mxFree(logpsbv_p2k);
		mxFree(tWZnt_dp2k);
	} else {
		const int nrvals = mxGetNumberOfElements(prhs[9]);
		double* const logpsrv_nrp1 = (double*)mxMalloc((nrvals+1)*sizeof(double));
		double* const tWZnt_d = (double*)mxMalloc(D*sizeof(double));

		if (VRBSLVL > 0) printf("[");

		for (int k = 0; k < K; ++k) {
			/* log-prior */
			logpsrv_nrp1[0] = 0;
			for (int i = 1; i <= nrvals; ++i) {
				logpsrv_nrp1[i] = logmbyNmm_k[k] - log((double)nrvals);
			}

			/* log-likelihood */
			for (int d = 0; d < D; ++d) {
				WZnt_d[d] -= W_dk_ptr[d]*Zn_k[k];
			}

			memset(tWZnt_d, 0, D*sizeof(double));
			INCLOGLKLHD(logpsrv_nrp1[0], tWZnt_d);
			
			for (int i = 0; i < nrvals; ++i) {
				for (int d = 0; d < D; ++d) {
					tWZnt_d[d] = rvals[i]*W_dk_ptr[d];
				}
				INCLOGLKLHD(logpsrv_nrp1[i+1], tWZnt_d);
			}

			if (VRBSLVL > 1) {
				for (int i = 0; i <= nrvals; ++i) {
					printf("%.2f ", logpsrv_nrp1[i]);
				}
			}

			int pid = from_logps(logpsrv_nrp1, nrvals+1, algtype);
			assert(pid >= 0);
			Zn_k[k] = pid==0 ? 0.0 : rvals[pid-1];
			
			for (int d = 0; d < D; ++d) {
				WZnt_d[d] += (*W_dk_ptr)*Zn_k[k];
				W_dk_ptr++;
			}

			if (VRBSLVL > 0) {
				printf("%.4f, ", Zn_k[k]);
			}
		}
		if (VRBSLVL > 0) printf("\b\b]\n");

		mxFree(logpsrv_nrp1);
		mxFree(tWZnt_d);
	}
}

int from_logps(double* const logps, const int N, const int algtype) {
	static std::mt19937 eng;
	static std::uniform_real_distribution<double> unif(0, 1);

	int maxpkid = 0;
	double maxlogpk = logps[0];
	for (int i = 1; i < N; ++i) {
		if (logps[i] > maxlogpk) {
			maxlogpk = logps[i];
			maxpkid = i;
		}
	}
	if (algtype == 1) {
		logps[0] = exp(logps[0]-maxlogpk);
		for (int i = 1; i < N; ++i) {
			logps[i] = logps[i-1] + exp(logps[i]-maxlogpk);
		}
		const double logsumpk = maxlogpk + log(logps[N-1]);
		const double logrand = log(unif(eng)) + logsumpk;
		for (int i = 0; i < N; ++i) {
			logps[i] = maxlogpk + log(logps[i]);
			if (logrand < logps[i]) {
				return i;
			}
		}
	} else {
		return maxpkid;
	}
	return -1;
}